class CondBlock(nn.Module):
def __init__(self, input_size, hidden_size, condition_size, last):
super().__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.condition_size = condition_size
self.split = self.input_size // 2
self.s_net = nn.Sequential(
nn.Linear(self.split + self.condition_size, self.hidden_size),
nn.ReLU(),
nn.Linear(self.hidden_size, self.hidden_size),
nn.ReLU(),
nn.Linear(self.hidden_size, self.split),
nn.Tanh(),
)
self.t_net = nn.Sequential(
nn.Linear(self.split + self.condition_size, self.hidden_size),
nn.ReLU(),
nn.Linear(self.hidden_size, self.hidden_size),
nn.ReLU(),
nn.Linear(self.hidden_size, self.split),
)
if last:
self.Q = torch.eye(self.input_size)
else:
self.Q, _ = torch.linalg.qr(torch.randn(size=(self.input_size, self.input_size)))
def encode(self, x, condition):
x1 = x[:, :self.split]
x2 = x[:, self.split:]
s = torch.exp(self.s_net(torch.cat((x1, condition), dim=1)))
t = self.t_net(torch.cat((x1, condition), dim=1))
x2 = s * x2 + t
return torch.cat((x1, x2), dim=1) @ self.Q
def decode(self, x, condition):
x = x @ self.Q.T
x1 = x[:, :self.split]
x2 = x[:, self.split:]
s = torch.exp(self.s_net(torch.cat((x1, condition), dim=1)))
t = self.t_net(torch.cat((x1, condition), dim=1))
x2 = (x2 - t) / s
return torch.cat((x1, x2), dim=1)
def loss_fwd(self, x, condition):
x1 = x[:, :self.split]
x2 = x[:, self.split:]
s_wig = self.s_net(torch.cat((x1, condition), dim=1))
s = torch.exp(s_wig)
t = self.t_net(torch.cat((x1, condition), dim=1))
x2 = s * x2 + t
return torch.cat((x1, x2), dim=1) @ self.Q, s_wig
class CondRealNVP(nn.Module):
def __init__(self, input_size, hidden_size, blocks, condition_size):
super().__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.condition_size = condition_size
self.blocks = blocks
self.couplings = nn.Sequential(*[
CondBlock(self.input_size, self.hidden_size, self.condition_size, i == self.blocks - 1)
for i in range(self.blocks)
])
self.sampler = MultivariateNormal(torch.zeros(self.input_size), torch.eye(self.input_size))
def encode(self, x, condition):
for b in self.couplings:
x = b.encode(x, condition)
return x
def decode(self, x, condition):
for b in self.couplings[::-1]:
x = b.decode(x, condition)
return x
def loss(self, x, condition):
s_vals = torch.zeros(len(x))
for b in self.couplings:
x, s_wig = b.loss_fwd(x, condition)
s_vals += torch.sum(s_wig, dim=1)
return torch.mean(torch.sum(x * x, dim=1) / 2 - s_vals)
def sample(self, n_samples, conditions):
with torch.no_grad():
generated = np.zeros((len(conditions), n_samples, self.input_size))
for i, condition in enumerate(conditions):
cond_stacked = torch.tensor([condition.tolist()] * n_samples).to(torch.float)
codes = self.sampler.sample((n_samples,))
generated[i] = self.decode(codes, cond_stacked)
return generated.reshape(len(conditions) * n_samples, self.input_size)
def print(self):
for name, param in self.named_parameters():
print(f"{name}, {param.shape}, grad: {param.requires_grad}")
def train_cinn(model, X_train, y_train, n_epochs, plot_loss=False, eval_data=None, data_loader=None):
optimizer = torch.optim.Adam(model.parameters())
if data_loader is None:
if eval_data is not None:
X_eval, y_eval = eval_data
eval_losses = []
losses = []
for i in tqdm(range(n_epochs)):
loss = model.loss(X_train, y_train)
loss.backward()
optimizer.step()
optimizer.zero_grad()
losses.append(loss.item())
if eval_data is not None:
model.eval()
with torch.no_grad():
eval_losses.append(model.loss(X_eval, y_eval).item())
model.train()
if plot_loss:
plt.figure(figsize=(12, 8))
if eval_data is not None:
plt.plot(range(n_epochs), losses, label="Train Loss")
plt.plot(range(n_epochs), eval_losses, label="Eval Loss")
plt.title("Loss during Training")
plt.legend()
else:
plt.plot(range(n_epochs), losses)
plt.title("Train Loss", size=18)
plt.xlabel("Epoch", size=16)
plt.ylabel("NLL", size=16)
plt.show()
else:
losses = []
for i in tqdm(range(n_epochs)):
epoch_loss = 0.0
for X, y in data_loader:
one_hot_y = nn.functional.one_hot(y.to(torch.int64), 10)
loss = model.loss(X.reshape(-1, 64), one_hot_y)
loss.backward()
optimizer.step()
optimizer.zero_grad()
epoch_loss += loss.item()
losses.append(epoch_loss)
if plot_loss:
plt.figure(figsize=(12, 8))
plt.plot(range(n_epochs), losses)
plt.title("Train Loss", size=18)
plt.xlabel("Epoch", size=16)
plt.ylabel("NLL", size=16)
plt.show()
